from io import StringIO

import pandas as pd
import psycopg2


def get_conn():
    connection = psycopg2.connect(
        database="net_platfrom",
        user="postgres",
        password="wxxx12345",
        host="192.168.67.56",
        port="5432"
    )
    return connection


def insert_csv_to_pg(csv_file, columns, table_name, encoding='utf-8'):
    # 连接到 PostgreSQL 数据库
    conn = get_conn()
    cursor = conn.cursor()

    # 给列名加上双引号
    if columns:
        quoted_columns = [col for col in columns]
        print(quoted_columns)
    else:
        quoted_columns = None
    full_table_name = table_name
    # 清空表
    print(f"Truncating table {full_table_name}...")
    cursor.execute(f'TRUNCATE TABLE "{full_table_name}";')
    conn.commit()
    # 逐块读取 CSV 文件
    chunksize = 10000
    for chunk in pd.read_csv(csv_file, chunksize=chunksize, dtype=str, encoding=encoding):
        # 选择指定列
        if columns:
            selected_chunk = chunk[columns]
        else:
            selected_chunk = chunk
        if quoted_columns is None:
            quoted_columns = chunk.columns.tolist()
        # 将 DataFrame 转换为 CSV 格式的字符串
        csv_buffer = StringIO()
        selected_chunk.to_csv(csv_buffer, sep='|', na_rep='nan', header=False, index=False)
        csv_buffer.seek(0)
        try:
            # 使用 COPY 命令插入数据
            cursor.copy_from(csv_buffer, full_table_name, sep='|', columns=quoted_columns)
            conn.commit()
            print(f"Inserted {len(selected_chunk)} rows.")
        except Exception as e:
            print(f"Error: {e}")
            conn.rollback()
    cursor.close()
    conn.close()


def copy_to_table(path, table):
    conn = get_conn()
    cur = conn.cursor()

    # 执行 COPY 命令
    copy_sql = f"COPY {table} FROM STDIN WITH (FORMAT CSV, HEADER)"
    with open(path, 'r', encoding='utf-8') as file:
        cur.copy_expert(copy_sql, file)

    # 提交事务
    conn.commit()

    # 关闭游标和连接
    cur.close()
    conn.close()
